[BugFix] honor collector masks in loss reductions - #4057
Open
gtnv wants to merge 7 commits into
Open
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4057
Note: Links to docs will display an error until the docs builds have been completed. ❌ 7 New FailuresAs of commit 42a022e with merge base 7330afb ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
gtnv
marked this pull request as ready for review
August 2, 2026 01:23
theap06
suggested changes
Aug 2, 2026
Collaborator
|
I made some edits in |
github-actions
Bot
added
the
Documentation
PR pytorch#3888 routed objectives through _reduce_loss, but the helper still ignored ("collector", "mask") because the corresponding change in pytorch#3850 never merged. This fix excludes padded timesteps from loss reduction so they do not affect the loss or gradients.
Unify unreduced mask behavior. Apply collector masks to REDQ under reduction="none". Remove unreachable mask-shape handling.
Keeps the pytorch#3866 design -- discovery of the validity masks TorchRL writes itself is the default, so switching to SliceSampler(pad_output=True) needs no per-loss wiring -- and adds the escape hatch it was missing, plus the coverage and the robustness fix. - `LossModule.loss_mask_key` (new, default `"auto"`) selects the mask source: `"auto"` reads every entry of `AUTO_LOSS_MASK_KEYS` and ANDs the ones it finds; a `NestedKey` restricts masking to that single entry, resolving which mask applies when a batch carries both; `None` disables masking. Exposed as an attribute rather than a constructor kwarg so it needs no matching field on every `*LossConfig` (CLAUDE.md section 14). - Masked positions are now selected out with `torch.where` on every reduction path, not multiplied by a zeroed mask: `nan * 0` is `nan` in both the forward and the backward pass, so a non-finite value at a masked position used to poison the loss and the gradients. - Tests: `TestLossMaskReduction` in test/objectives/test_loss_module.py pins the three mask-source modes, the AND of the auto keys, no-op reduction on unmasked data, non-finite masked positions, trailing-dim broadcasting, and both `reduction="none"` contracts. `test_collector_mask` / `test_collector_mask_can_be_disabled` in test/objectives/test_bc.py cover a real loss end to end. - Docs: a "Masked reduction" section in objectives_common.rst, and `AUTO_LOSS_MASK_KEYS` exported from `torchrl.objectives`. Reverts the unannounced `BCLoss` change that turned an all-padding batch from 0 into `nan`: worth doing, but it is a separate behavior change that needs its own deprecation note and test. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Running the objectives suite locally shows the collector-mask discovery breaks
72 existing tests: `("collector", "mask")` is [B, T, 1] (masks carry a trailing
singleton to broadcast against [..., 1]-shaped rewards) while a per-timestep
loss is [B, T], and _expand_loss_mask only handled masks with *fewer* dims than
the loss, so expand_as raised. TD3/TD3+BC batcher tests cover exactly this
shape.
_expand_loss_mask now drops trailing singleton dimensions before broadcasting,
and raises a clear ValueError when the mask has genuinely more non-singleton
dimensions than the loss. Regression tests for both, parametrized over the two
auto keys.
test/objectives: 7522 passed, matching the pre-PR baseline exactly.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
vmoens
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
#3888 refactored objectives to route through _reduce_loss but the collector mask change from #3850 never merged, leaving a gap. The relevant code changed as well. This made it so padded timesteps from collector, mask were still included and could change the final loss and grads.
BC & RND use it directly. A2C, ppo and reinforce pass input via TensorDict and REDQ/OnlineDT deal with shape-specific cases their own way.
closes #3866.
Ran tests; 9314 passed.
cc @vmoens
Tested
test/objectives: 7525 passed, 2330 skipped, matching the pre-PR baseline(7522 passed with the same skips, plus the 3 new regression tests).
Worth flagging: before the fix in
4cc1e15, this branch failed 72 existingtests — every
TestTD3::test_td3_batcherandTestTD3BC::test_td3bc_batcherparametrization. Their fixtures supply
("collector", "mask")with the standardtrailing-singleton shape
[B, T, 1]against a[B, T]loss, and_expand_loss_maskonly handled masks with fewer dimensions than the loss, soexpand_asraised.shifted_validnever hit this because the value estimatorsemit it already squeezed. That crash is the reason a mask-discovery change needs
the objectives suite run against a baseline, not just new unit tests.